import torch
from torch import nn
class Predictor(nn.Module):

    def __init__(self):
        super().__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(4, 50, 3, padding=1, bias=True),
            nn.LayerNorm(128),
            #nn.LeakyReLU(negative_slope=0.125),
            nn.Conv2d(50, 50, 3, padding=1, bias=True),
            nn.LeakyReLU(negative_slope=0.125),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(50, 50, 3, padding=1, bias=True),
            nn.LayerNorm(128),
            nn.Conv2d(50, 50, 3, padding=1, bias=True),
            nn.LeakyReLU(negative_slope=0.125),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(50, 64, 3, padding=1, bias=True),
            nn.LayerNorm(128),
            nn.Conv2d(64, 4, 3, padding=1, bias=True),
            nn.LeakyReLU(negative_slope=0.125),
        )

    def forward(self, x):
        n, c, h, w = x.shape
        x = x.reshape((n, c, h // 2, 2, w // 2, 2)).permute((0, 1, 3, 5, 2, 4)).reshape((n, c * 4, h // 2, w // 2))
        x = self.conv1(x)
        residential = x
        x = self.conv2(x)
        x += residential
        x = self.conv3(x)
        x = x.reshape((n, c, 2, 2, h // 2, w // 2)).permute((0, 1, 4, 2, 5, 3)).reshape((n, c, h, w))
        return x

if __name__ == '__main__':
    x= torch.randn(2,1,256,256)
    net = Predictor()
    # n, c, h, w = x.shape
    # x = x.reshape((n, c, h // 2, 2, w // 2, 2)).permute((0, 1, 3, 5, 2, 4)).reshape((n, c * 4, h // 2, w // 2))
    total = sum([param.nelement() for param in net.parameters()])
    print("Number of parameter: %.2fM" % (total / 1e6))
    #print(x.shape)

    print(net)
    print(net(x).shape)